% Basic Introduction:
% This function computes the transition prob. matrix from an initial grid X
% to a given a fixed histgram approximation grid HistGrid. In this function,
% Xp is the transition destination of X. 
% Function form:
% [varargout]=AssignGridTrProb(HistGrid,Xp,TypeFlag,OutputFlag)
% Input Variables:
%       HistGrid            Histgram grid 
%       Xp                  Transition destination from X
%       TypeFlag            Flag for the type of assignment
%                           'Continuous': assign the prob. according to the
%                           distance between destination interval end
%                           points and the destination points. 
%                           'Half-Half': assign the prob. to both
%                           destination interval end points with 1/2
%       OutputFlag          Flag for the type of output
%                           'Index & TrProb': output the index and the
%                           corresponding transition prob. value
%                           'TrProbMat': a sparse matrix for transition
%                           prob.
% Special Notes:
% This function use the external Mex function lookup from CompEcon toolbox.

function [varargout]=AssignGridTrProb(HistGrid,Xp,TypeFlag,OutputFlag)
L_Grid      =   length(HistGrid);
L_X         =   length(Xp);
Ind         =   lookup(HistGrid,Xp);

I_leftout   =   (Ind==0);
I_rightout  =   (Ind==L_Grid);
% If the destination is to the left of left histgram boundary, 
% its destination interval left end will be assigned as the left boundary
% Ind_left    =   max(Ind,1);
Ind_left    =   Ind;
Ind_left(I_leftout)     =   1;
% If the destination is to the right of right histgram boundary, its
% destination interval right end will be assigned as the right boundary
% Ind_right   =   min(Ind+1,L_Grid);
Ind_right  =   Ind+1;
Ind_right(I_rightout)   =   L_Grid;
% For those with destination point outside of boundary, their left/right
% destination interval end points are the same.

switch TypeFlag
    case 'Continuous'
        % In this case, the transition prob. to the interval end points are
        % assigned proportional to their distance to the destination point.
        Gap_left    =   Xp-HistGrid(Ind_left);
        Gap_right   =   HistGrid(Ind_right)-Xp;
        TrProb_left =   Gap_left./(Gap_left+Gap_right);
        TrProb_left(I_leftout)      =   0;
        TrProb_left(I_rightout)     =   1;
        TrProb_right=   1-TrProb_left;
            
    case 'Half-Half'
        % In this case, the transition prob. to the interval end points are
        % assigned both as 1/2.
        TrProb_left =   repmat(1/2,[L_X,1]);
        TrProb_left(I_leftout)      =   0;
        TrProb_left(I_rightout)     =   1;
        TrProb_right=   1-TrProb_left;
end

XpColInd            =   (1:1:L_X)';
TrProbMatRowInd     =   [Ind_left(~I_leftout);Ind_right(~I_rightout)];
TrProbMatColInd     =   [XpColInd(~I_leftout);XpColInd(~I_rightout)];
TrProbMatValue      =   [TrProb_left(~I_leftout);TrProb_right(~I_rightout)];
switch OutputFlag
    case 'Index & TrProb'
        varargout{1}=   [TrProbMatRowInd,TrProbMatColInd];
        varargout{2}=   TrProbMatValue;
    case 'TrProbMat'
        varargout{1}=   sparse(TrProbMatRowInd,TrProbMatColInd,TrProbMatValue,...
                               L_Grid,L_X);
end


